import socket
import http.server
import json
from qkdKey import QKDKey
from qkdBuffer import QKDBuffer
import threading

class QKDKeyManagerSystem:
    def __init__(self, kms_id, max_keys_per_request, min_key_size, max_key_size, kms_hostname):
        self.kms_id = kms_id  # Unique identifier for the Key Management System (KMS)
        self.max_keys_per_request = max_keys_per_request  # Maximum number of keys that can be supplied in a single request
        self.min_key_size = min_key_size  # Minimal key size allowed to request
        self.max_key_size = max_key_size  # Maximal key size allowed to request
        self.kms_hostname = kms_hostname  # Unique hostname for the KMS
        self.buffers = {}  # Dictionary to store multiple QKDBuffers identified by their unique IP address

    class RequestHandler(http.server.BaseHTTPRequestHandler):
        def do_GET(self):
            path_parts = self.path.split('/')
            print(f"path_parts:{path_parts}")
            if len(path_parts) > 3 and path_parts[1] == 'api':
                if path_parts[2] == "keys" and path_parts[3] == "status":
                    self.send_response(200)
                    self.send_header("Content-type", "application/json")
                    self.end_headers()
                    status = {
                        "kms_id": self.server.kms.kms_id,
                        "buffers_status": self.server.kms.get_status()
                    }
                    self.wfile.write(bytes(json.dumps(status), "utf8"))
                elif path_parts[2] == "keys" and "request_keys" in path_parts[3]:
                    try:
                        query_components = self.path.split("?")[1].split("&")
                        print(f"query_components:{query_components}")
                        query_dict = {key: value for key, value in (item.split("=") for item in query_components)}
                        num_keys = int(query_dict.get("num_keys"))
                        print(f"tag1 num_keys:{num_keys}")
                        ip_address = query_dict.get("ip_address")
                        print(f"tag2 ip_address:{ip_address}")
                        supplied_keys = self.server.kms.request_keys(num_keys, ip_address)
                        print(f"supplied_keys:{supplied_keys}")
                        self.send_response(200)
                        self.send_header("Content-type", "application/json")
                        self.end_headers()
                        keys_info = [{"key_id": key.key_id, "key_value": key.key_value.hex()} for key in supplied_keys]
                        print(f"keys_info:{keys_info}")
                        self.wfile.write(json.dumps(keys_info).encode("utf8"))
                    except Exception as e:
                        self.send_response(400)
                        self.send_header("Content-type", "text/plain")
                        self.end_headers()
                        self.wfile.write(bytes(f"Error: {str(e)}", "utf8"))
                elif path_parts[2] == "keys" and "get_key_with_ids" in path_parts[3]:
                    try:
                        query_components = self.path.split("?")[1].split("&")
                        print(f"query_components:{query_components}")
                        query_dict = {key: value for key, value in (item.split("=") for item in query_components)}
                        print(f"query_dict:{query_dict}")
                        ip_address = query_dict.get("ip_address")
                        key_ids = [int(key_id) for key_id in query_dict.get("key_IDs", "").split(",")]
                        print(f"Getting keys by IDs for IP address: {ip_address}, Key IDs: {key_ids}")
                        supplied_keys = self.server.kms.get_keys_by_ids(ip_address, key_ids)
                        self.send_response(200)
                        self.send_header("Content-type", "application/json")
                        self.end_headers()
                        keys_info = [{"key_id": key.key_id, "key_value": key.key_value.hex()} for key in supplied_keys]
                        self.wfile.write(json.dumps(keys_info).encode("utf8"))
                    except Exception as e:
                        self.send_response(400)
                        self.send_header("Content-type", "text/plain")
                        self.end_headers()
                        self.wfile.write(bytes(f"Error: {str(e)}", "utf8"))
                elif path_parts[2] == "keys" and "query_keys_for_ip" in path_parts[3]:
                    try:
                        query_components = self.path.split("?")[1].split("&")
                        query_dict = {key: value for key, value in (item.split("=") for item in query_components)}
                        ip_address = query_dict.get("ip_address")
                        print(f"Getting all key IDs for IP address: {ip_address}")
                        if ip_address not in self.server.kms.buffers:
                            raise ValueError(f"No buffer found with IP address {ip_address}")
                        buffer = self.server.kms.buffers[ip_address]
                        key_ids = list(buffer.buffer.keys())
                        self.send_response(200)
                        self.send_header("Content-type", "application/json")
                        self.end_headers()
                        self.wfile.write(bytes(json.dumps({"key_ids": key_ids}), "utf8"))
                    except Exception as e:
                        print(f"Error while handling GET /keys/query_keys_for_ip: {e}")
                        self.send_response(502)
                        self.send_header("Content-type", "text/plain")
                        self.end_headers()
                        self.wfile.write(bytes(f"Bad Gateway (502): {str(e)}", "utf8"))
                else:
                    self.send_response(404)
                    self.send_header("Content-type", "text/plain")
                    self.end_headers()
                    self.wfile.write(bytes("Not api1 Found", "utf8"))
            else:
                self.send_response(404)
                self.send_header("Content-type", "text/plain")
                self.end_headers()
                self.wfile.write(bytes("Not api2 Found", "utf8"))

        def do_POST(self):
            path_parts = self.path.split('/')
            if len(path_parts) > 3 and path_parts[1] == 'api':
                if path_parts[2] == "keys" and path_parts[3] == "add_key":
                    try:
                        content_length = int(self.headers['Content-Length'])
                        print(f"tag2 content_length:{content_length}")
                        post_data = self.rfile.read(content_length)
                        request_data = json.loads(post_data)
                        ip_address = request_data.get("ip_address")
                        key_value = request_data.get("key_value")
                        key_id = int(request_data.get("key_id"))  # Ensure key_id is an int
                        print(f"tag3 ip_address:{ip_address}")
                        print(f"tag4 key_value:{key_value}")
                        print(f"tag5 key_id:{key_id}")
                        key = QKDKey(key_value=bytes.fromhex(key_value), key_id=key_id)
                        self.server.kms.add_key_to_buffer(ip_address, key)
                        self.send_response(200)
                        self.send_header("Content-type", "text/plain")
                        self.end_headers()
                        self.wfile.write(bytes(f"Key added to buffer with IP address {ip_address}.", "utf8"))
                    except Exception as e:
                        self.send_response(400)
                        self.send_header("Content-type", "text/plain")
                        self.end_headers()
                        self.wfile.write(bytes(f"Error: {str(e)}", "utf8"))
                else:
                    self.send_response(404)
                    self.send_header("Content-type", "text/plain")
                    self.end_headers()
                    self.wfile.write(bytes("Not Found", "utf8"))
            else:
                self.send_response(404)
                self.send_header("Content-type", "text/plain")
                self.end_headers()
                self.wfile.write(bytes("Not Found", "utf8"))

    def add_buffer(self, buffer):
        if not isinstance(buffer, QKDBuffer):
            print("Invalid buffer type. Only QKDBuffer instances can be added.")
            return
        if buffer.ip_address in self.buffers:
            print(f"A buffer with IP address {buffer.ip_address} already exists.")
        else:
            self.buffers[buffer.ip_address] = buffer
            print(f"Buffer with IP address {buffer.ip_address} added to KMS.")

    def remove_buffer(self, ip_address):
        if ip_address in self.buffers:
            del self.buffers[ip_address]
            print(f"Buffer with IP address {ip_address} removed from KMS.")
        else:
            print(f"No buffer found with IP address {ip_address}.")

    def add_key_to_buffer(self, ip_address, key):
        if ip_address not in self.buffers:
            print(f"No buffer found with IP address {ip_address}.")
            return
        buffer = self.buffers[ip_address]
        buffer.add_key(key)
        print(f"Key added to buffer with IP address {ip_address}. key ID:{key.key_id}.key value:{key.key_value.hex()}")

    def request_keys(self, num_keys, ip_address):
        if ip_address not in self.buffers:
            print(f"No buffer found with IP address {ip_address}.")
            return []
        buffer = self.buffers[ip_address]
        if num_keys > self.max_keys_per_request:
            print(f"Request exceeds maximum keys per request limit: {self.max_keys_per_request}")
            return []

        supplied_keys = []
        for _ in range(num_keys):
            if len(buffer.buffer) == 0:
                print("Buffer is empty, no more keys to supply.")
                break
            key_id = list(buffer.buffer.keys())[0]
            key = buffer.get_remove_key(key_id)
            if key:
                supplied_keys.append(key)

        print(f"Supplied {len(supplied_keys)} keys from buffer with IP address {ip_address}.")
        return supplied_keys

    def get_keys_by_ids(self, ip_address, key_ids):
        if ip_address not in self.buffers:
            print(f"No buffer found with IP address {ip_address}.")
            return []
        buffer = self.buffers[ip_address]
        supplied_keys = []
        for key_id in key_ids:
            key = buffer.get_remove_key(key_id)
            if key:
                supplied_keys.append(key)
            else:
                print(f"Key with ID {key_id} not found in buffer with IP address {ip_address}.")

        print(f"Supplied {len(supplied_keys)} keys by IDs from buffer with IP address {ip_address}.")
        return supplied_keys

    def get_status(self):
        status = {
            ip_address: {
                "buffer_status": buffer.get_status(),
                "key_count": len(buffer.buffer)
            } for ip_address, buffer in self.buffers.items()
        }
        print(f"KMS Buffers Status: {status}")
        return status

    # def run_server(self):
    #     print(f"KMS server running on all associated buffer IP addresses...")
    #     server_addresses = [(buffer.ip_address, 80) for buffer in self.buffers.values()]
    #     httpd_list = [http.server.HTTPServer(address, self.RequestHandler) for address in server_addresses]
    #     for server in httpd_list:
    #         server.kms = self  # Attach KMS instance to the HTTP server
    #         server.serve_forever()
    
    def run_server(self):
        print(f"KMS server running on all associated buffer IP addresses...")
        server_addresses = [(buffer.ip_address, 8080) for buffer in self.buffers.values()]
        httpd_list = [http.server.HTTPServer(address, self.RequestHandler) for address in server_addresses]
        # httpd_list = [http.server.HTTPServer(address, http.server.SimpleHTTPRequestHandler) for address in server_addresses]
        for server in httpd_list:
            try:
                print(f"Attempting to start server on {server.server_address}")
                server.kms = self  # Attach KMS instance to the HTTP server
                threading.Thread(target=server.serve_forever, daemon=False).start()
                print(f"Successfully started server on {server.server_address}")
            except Exception as e:
                print(f"Failed to start server on {server.server_address}: {e}")

if __name__ == "__main__":
    kms = QKDKeyManagerSystem(kms_id="KMS001", max_keys_per_request=5, min_key_size=128, max_key_size=256, kms_hostname="kms001.example.com")
    buffer1 = QKDBuffer(key_size=256, max_capacity=200, min_capacity=20, ip_address="192.168.1.1")
    buffer2 = QKDBuffer(key_size=256, max_capacity=100, min_capacity=10, ip_address="192.168.1.2")
    # buffer1 = QKDBuffer(key_size=256, max_capacity=200, min_capacity=20, ip_address="127.0.0.1")
    # buffer2 = QKDBuffer(key_size=256, max_capacity=100, min_capacity=10, ip_address="127.0.0.1")
# 
    kms.add_buffer(buffer1)
    kms.add_buffer(buffer2)
    # Example key management process
    kms.get_status()
    # Run the HTTP server to manage key requests and status inquiries
    kms.run_server()
